"""
Demo 04: SAX符号化
将时间序列转换为符号表示，便于离散化处理和模式匹配
"""

import numpy as np
from demo_01_raw import RawVectorizer
try:
    from tslearn.piecewise import SymbolicAggregateApproximation
except ImportError:
    print("请先安装tslearn: pip install tslearn")
    raise


class SAXVectorizer:
    """SAX符号化向量处理类"""
    
    def __init__(self, window_size=60, step_size=5, n_bins=8, alphabet='abcdefgh'):
        """
        初始化参数
        
        Args:
            window_size: 窗口大小
            step_size: 步长
            n_bins: 分箱数（字母表大小）
            alphabet: 使用的字母表
        """
        self.window_size = window_size
        self.step_size = step_size
        self.n_bins = n_bins
        self.alphabet = alphabet
        self.raw_vectorizer = RawVectorizer(window_size, step_size)
        self.sax = SymbolicAggregateApproximation(
            n_segments=window_size,  # 段数等于窗口长度
            alphabet_size_avg=n_bins
        )
        self.char_to_int = {c: i for i, c in enumerate(alphabet)}
        
    def extract_sax_features(self, windows):
        """
        将窗口转换为SAX符号表示
        
        Args:
            windows: 窗口矩阵 (n_windows, window_size)
            
        Returns:
            numpy.ndarray: SAX向量矩阵 (n_windows, window_size)
        """
        # 转换为SAX符号
        X_sax = self.sax.fit_transform(windows)
        
        # 将字符映射为整数
        X_sax_int = np.vectorize(self.char_to_int.get)(X_sax).astype(float)
        
        # Reshape为2D矩阵
        X_sax_vec = X_sax_int.reshape(X_sax_int.shape[0], -1)
        
        return X_sax_vec
    
    def extract_sax_histogram(self, windows):
        """
        提取SAX符号的直方图特征
        
        Args:
            windows: 窗口矩阵
            
        Returns:
            numpy.ndarray: 符号直方图特征矩阵
        """
        # 获取SAX符号表示
        X_sax = self.sax.fit_transform(windows)
        
        # 计算每个窗口的符号直方图
        histograms = []
        for window_sax in X_sax:
            hist = np.zeros(self.n_bins)
            for char in window_sax.flatten():
                if char in self.char_to_int:
                    hist[self.char_to_int[char]] += 1
            hist = hist / len(window_sax.flatten())  # 归一化
            histograms.append(hist)
        
        return np.array(histograms)
    
    def extract_sax_ngrams(self, windows, n=2):
        """
        提取SAX n-gram特征
        
        Args:
            windows: 窗口矩阵
            n: n-gram的n值
            
        Returns:
            dict: n-gram频率字典
        """
        # 获取SAX符号表示
        X_sax = self.sax.fit_transform(windows)
        
        all_ngrams = {}
        for window_sax in X_sax:
            window_str = ''.join(window_sax.flatten())
            for i in range(len(window_str) - n + 1):
                ngram = window_str[i:i+n]
                all_ngrams[ngram] = all_ngrams.get(ngram, 0) + 1
        
        return all_ngrams
    
    def compute_sax_distance(self, sax1, sax2):
        """
        计算两个SAX表示之间的距离
        
        Args:
            sax1: 第一个SAX向量
            sax2: 第二个SAX向量
            
        Returns:
            float: 距离值
        """
        return np.sqrt(np.sum((sax1 - sax2) ** 2))
    
    def fit_transform(self, ticker="AAPL", start="2020-01-01", end="2024-12-31"):
        """
        完整的SAX向量化流程
        
        Args:
            ticker: 股票代码
            start: 开始日期
            end: 结束日期
            
        Returns:
            tuple: (原始窗口矩阵, SAX向量矩阵, SAX直方图矩阵)
        """
        # 获取原始窗口
        X_raw = self.raw_vectorizer.fit_transform(ticker, start, end)
        
        # 提取SAX特征
        X_sax = self.extract_sax_features(X_raw)
        
        # 提取SAX直方图
        X_sax_hist = self.extract_sax_histogram(X_raw)
        
        return X_raw, X_sax, X_sax_hist
    
    def decode_sax(self, sax_vector):
        """
        将SAX整数向量解码回字符串
        
        Args:
            sax_vector: SAX整数向量
            
        Returns:
            str: SAX字符串表示
        """
        int_to_char = {v: k for k, v in self.char_to_int.items()}
        chars = [int_to_char.get(int(x), '?') for x in sax_vector]
        return ''.join(chars)


def main():
    """主函数，演示使用方法"""
    print("=" * 60)
    print("Demo 04: SAX符号化")
    print("=" * 60)
    
    # 创建SAX向量化器
    vectorizer = SAXVectorizer(window_size=60, step_size=5, n_bins=8)
    
    # 执行向量化
    print("正在下载AAPL股票数据并进行SAX符号化...")
    X_raw, X_sax, X_sax_hist = vectorizer.fit_transform(
        ticker="AAPL",
        start="2020-01-01",
        end="2024-12-31"
    )
    
    # 输出结果
    print(f"\n原始窗口 shape: {X_raw.shape}")
    print(f"SAX向量 shape: {X_sax.shape}")
    print(f"SAX直方图 shape: {X_sax_hist.shape}")
    
    # 显示第一个窗口的SAX表示
    print(f"\n第一个窗口的SAX表示（前20个符号）:")
    sax_string = vectorizer.decode_sax(X_sax[0][:20])
    print(f"  符号串: {sax_string}")
    print(f"  整数表示: {X_sax[0][:20].astype(int)}")
    
    # 显示符号分布
    print(f"\n第一个窗口的符号分布:")
    for i, char in enumerate(vectorizer.alphabet):
        print(f"  {char}: {X_sax_hist[0, i]:.3f}")
    
    # 计算窗口间的SAX距离
    dist_01 = vectorizer.compute_sax_distance(X_sax[0], X_sax[1])
    dist_02 = vectorizer.compute_sax_distance(X_sax[0], X_sax[2])
    print(f"\n窗口间的SAX距离:")
    print(f"  窗口0-窗口1: {dist_01:.4f}")
    print(f"  窗口0-窗口2: {dist_02:.4f}")
    
    # 提取2-gram特征
    ngrams = vectorizer.extract_sax_ngrams(X_raw[:10], n=2)
    top_ngrams = sorted(ngrams.items(), key=lambda x: x[1], reverse=True)[:5]
    print(f"\n前5个最频繁的2-gram:")
    for ngram, count in top_ngrams:
        print(f"  {ngram}: {count}次")
    
    return X_raw, X_sax, X_sax_hist


if __name__ == "__main__":
    X_raw, X_sax, X_sax_hist = main()
